import asyncio
import json

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from starlette.staticfiles import StaticFiles

from xhs_auto_poster.__001__langgraph_more_nodes.langgraph_more_node import agent_response
from xhs_auto_poster.__002__fastapi.msg_queue import msg_queue, MSG_TYPE
from xhs_auto_poster.common import project_path
from xhs_auto_poster.common.logger_util import get_logger

logger = get_logger(__name__)

app = FastAPI()

app.mount("/picture", StaticFiles(directory=project_path.joinpath("picture")))


async def event_generator(input):
    # 异步执行 Graph
    async def run_graph(input):
        try:
            logger.debug(f"在后台任务中开始执行Graph处理")
            output = await agent_response(input)
            await msg_queue.put(
                {"type": MSG_TYPE.DONE.name, "content": output})  # 结束标识
            logger.debug("Graph处理完成并发送结束标识")
        except Exception as e:
            logger.error(f"Graph处理任务出错: {str(e)}", exc_info=True)
            await msg_queue.put(
                {"type": MSG_TYPE.DONE.name, "content": f"处理出错: {str(e)}"})

    asyncio.create_task(run_graph(input))
    logger.info(f"创建Graph处理异步任务，开始监听消息队列")

    while True:
        msg = await msg_queue.get()
        logger.debug(f"从消息队列获取到消息: {msg.get('type')}")

        # 推送 SSE
        yield f"data: {json.dumps(msg, ensure_ascii=False)}\n\n"
        if msg.get("type") == MSG_TYPE.DONE.name:
            logger.debug("收到结束标识，退出事件生成器")
            break


@app.post("/process")
async def process(request: Request):
    # 获取传入的 JSON 数据
    try:
        logger.info("收到/process请求")
        data = await request.json()
        input = data.get("input", "")
        logger.info(f"请求处理中，输入内容: {input[:50]}..." if len(input) > 50 else f"请求处理中，输入内容: {input}")
        return StreamingResponse(event_generator(input), media_type="text/event-stream")
    except Exception as e:
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    import uvicorn
    
    logger.info("准备启动FastAPI服务")
    uvicorn.run(app, host="0.0.0.0", port=8001)
